import argparse
import cv2
import numpy as np
import matplotlib.pyplot as plt

# 参数解析
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", required=True, help="path to input image")
ap.add_argument("-t", "--template", required=True, help="path to template image")
args = vars(ap.parse_args())


# 图像显示函数，全文唯一的自定义函数
def cv_show(name, img):
    cv2.imshow(name, img)  # (自定义图像名,图像变量)
    cv2.waitKey(0)  # 图像窗口不会自动关闭
    cv2.destroyAllWindows()  # 手动关闭窗口


# 读取模板图像
reference = cv2.imread(args['template'])  # 获取指定文件夹下的某张图片

# 转换灰度图
ref = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY)

# 二值化处理
ret, thresh = cv2.threshold(ref, 127, 255, cv2.THRESH_BINARY_INV)

# 轮廓检测
contours, hierarchy = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 绘制轮廓
draw = reference.copy()
res = cv2.drawContours(draw, contours, -1, (0, 0, 255), 2)

# 模板排序
boxing = [cv2.boundingRect(cnt) for cnt in contours]
contours = list(contours)
contours, boxing = zip(*(sorted(zip(contours, boxing), key=lambda b: b[1][0])))

# 遍历每一个轮廓，给每一个轮廓对应具体数字
digits = {}
plt.figure(figsize=(10, 4))
for i, c in enumerate(contours):  # 返回轮廓下标和对应的轮廓值
    (x, y, w, h) = cv2.boundingRect(c)
    roi = ref[y:y + h, x:x + w]
    roi = cv2.resize(roi, (57, 88))

    # cv_show(f'{roi}', roi)

    plt.subplot(2, 5, i + 1)
    plt.title(i)  # 数字的序号
    plt.imshow(roi, 'gray')
    plt.xticks([])
    plt.yticks([])

    # 每一个数字对应的一个模板
    digits[i] = roi

# plt.show()

# 初始化卷积核
rectKernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 3))

img = cv2.imread(args['image'])
img = cv2.resize(img, (300, 200))
cv_show('img', img)

img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
cv_show('img_gray', img_gray)

# 礼帽操作 突出更明亮的区域
tophat = cv2.morphologyEx(img_gray, cv2.MORPH_TOPHAT, rectKernel)
cv_show('tophat', tophat)

gradX = cv2.Sobel(tophat, ddepth=cv2.CV_32F, dx=1, dy=0, ksize=-1)

gradX = np.absolute(gradX)
(minVal, maxVal) = (np.min(gradX), np.max(gradX))
gradX = (255 * ((gradX - minVal) / (maxVal - minVal)))
gradX = gradX.astype("uint8")

print(np.array(gradX).shape)
cv_show('gradX', gradX)

# 通过闭操作 (先膨胀再腐蚀) 将数字连在一起
gradX = cv2.morphologyEx(gradX, cv2.MORPH_CLOSE, rectKernel)
cv_show('gradX', gradX)

# cv2.THRESH_OTSU 会自动寻找合适的阈值, 适合双峰, 需要把阈值参数设置为0
thresh = cv2.threshold(gradX, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
cv_show('thresh', thresh)

# 再来一个闭操作
thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, rectKernel)
cv_show('thresh', thresh)